# -*- coding: UTF-8 -*-

# %% [markdown]
import os
import re
import numpy as np
import matplotlib.pyplot as plt
from scipy.optimize import curve_fit
import sys
import sympy as sp


def fourier_series(coeffs):
    theta = sp.symbols('theta')
    return np.sum([coeff[0]*sp.cos(coeff[1]*theta+coeff[2]) for coeff in coeffs])


def my_fft(Ts, y: np.ndarray):
    sp = np.fft.fft(y)
    freq = np.fft.fftfreq(y.shape[-1])
    amplitude = (abs(sp)/len(y)*2)[:int(len(freq)/2)]
    freq_Hz = (freq/Ts)[:int(len(freq)/2)]
    phase = np.arctan(abs(sp.imag)/abs(sp.real))[:int(len(freq)/2)]

    return freq_Hz, amplitude, phase


if __name__ == "__main__":
    pass

    # %% [markdown]
    # 模拟真实编码器数据
    encoder_max_pulse = 16384
    rad2encoder = encoder_max_pulse/(2*np.pi)
    encoder2rad = (2*np.pi)/encoder_max_pulse

    num = np.arange(0, 200)
    theta_idel = num/200*2*np.pi
    # plt.scatter(num, theta_idel)

    theta_ripple = 0.2*np.cos(2*theta_idel) + \
        0.08*np.cos(4*theta_idel + np.pi/2)
    #     0.2*np.cos(8*theta_idel) + \
    #     0.1*np.cos(14*theta_idel)

    theta_real = theta_idel - theta_ripple

    encoder_real = theta_real / (2*np.pi)*2**14

    theta_err = theta_idel - theta_real
    plt.figure(0)
    plt.plot(num, theta_ripple, marker='.')
    plt.figure(1)
    plt.plot(theta_idel, theta_err, marker='.')
    plt.figure(2)
    plt.plot(theta_real, theta_err, marker='.')
    # plt.plot(num, encoder_real, marker='.')

    plt.show()

    # %% [markdown]
    # 读取真实数据
    with open('./data.txt', 'r') as f:
        line = f.readline().splitlines()

        i = 0
        data_list = []
        while line:
            i += 1
            data = str(*line).split(' ')
            data = [float(x) for x in data]
            data_list.append(data)
            # print ("{}: {} | {}".format(i, data, type(data)))
            line = f.readline().splitlines()
        print("size:{}".format(len(data_list)))
    data_list = np.array(data_list)
    theta_idel = np.array(data_list[:, 0]*encoder2rad)
    theta_real = np.array(data_list[:, 1]*encoder2rad)
    theta_err = theta_idel - theta_real
    plt.figure(0)
    plt.plot(theta_idel, marker='.')
    plt.figure(1)
    plt.plot(theta_err, marker='.')
    plt.figure(2)
    plt.plot(theta_real, marker='.')
    # %% [markdown]
    # fft分析,取主导成分重构函数
    # plt.scatter(num, theta_ripple)
    freq, amp, phase = my_fft(1/200, theta_err)

    # plt.plot(freq, amp)
    # plt.xlim([0, 30])
    # plt.show()

    # plt.plot(freq, phase)
    # plt.xlim([0, 30])
    # plt.show()

    amp_threshold = 0.01
    main_freq = freq[amp >= amp_threshold]
    main_amp = amp[amp >= amp_threshold]
    main_phase = phase[amp >= amp_threshold]

    main_coeffs = np.array([main_amp, main_freq, main_phase]).T
    print(main_coeffs)
    f_expr = fourier_series(main_coeffs)

    num = np.arange(0, 200)
    theta_data = num/200*2*np.pi
    theta = sp.symbols('theta')

    f_data = [float(f_expr.evalf(subs={theta: data})) for data in theta_data]
    plt.plot(theta_data, f_data, marker='.')

    # plt.xlim([0,1000])
    # plt.show()
